# Python Standard Library Imports
import os
import time
from pathlib import Path as PathlibPath
from os.path import isfile, join

# Third-Party Imports
import numpy as np
from PIL import Image
import imageio
import argparse


def generate_gif_from_images(PNGS_DIR, GIF_FILE_PATH, IMAGE_TEMPLATE):
    # 记录gif生成时间,用于评估效率,没有特殊用途
    tic = time.time()
    
    # 创建一个图像列表,用于存储所有图像
    image_list = []
    
    # 获取目录下PNG图像文件的列表
    if not os.path.exists(PNGS_DIR):
        print(f"路径 {PNGS_DIR} 不存在!")
    
    files_in_directory = [file for file in PNGS_DIR.iterdir() if file.is_file() and file.name.endswith(".png")]
    
    # 找到文件夹中存在的最大编号
    max_image_index = -1
    for file in files_in_directory:
        try:
            index = int(file.stem.split('_')[-1])
            max_image_index = max(max_image_index,index)
        except ValueError:
            pass
    
    # 循环读取图片并添加到列表中
    for i in range(max_image_index + 1):
        image_path = PNGS_DIR / IMAGE_TEMPLATE.format(i)
        if image_path.exists():
            img = Image.open(image_path)
            image_list.append(img)
    
    # 保存GIF动画
    if image_list:
        image_list[0].save(GIF_FILE_PATH,save_all=True,append_images=image_list[1:],duration=100,loop=0)  # duration是指每一帧在gif中的显示时间（停留时间）
    
    print(f"###log### GIF文件保存在 {GIF_FILE_PATH}")
    toc = time.time()
    print(f"###log### GIF generated in {toc - tic:.2f} seconds.")


def generate_gif_for_single_folder(folder_path):
    """
    为单个文件夹生成GIF动画
    
    Args:
        folder_path (str or Path): 包含PNG图像的文件夹路径
    """
    folder_path = PathlibPath(folder_path)
    if not folder_path.exists():
        print(f"路径 {folder_path} 不存在!")
        return
        
    PNGS_DIR = folder_path / "images"
    GIF_FILE_PATH = folder_path / f"{folder_path.stem}.gif"
    IMAGE_TEMPLATE = f'{folder_path.stem}_{{}}.png'
    
    generate_gif_from_images(PNGS_DIR, GIF_FILE_PATH, IMAGE_TEMPLATE)


if __name__ == "__main__":
    # 创建命令行参数解析器
    parser = argparse.ArgumentParser(description='生成GIF动画')
    parser.add_argument('--folder', type=str, help='指定要生成GIF的文件夹路径')
    args = parser.parse_args()
    
    # 如果指定了文件夹，则只生成该文件夹的GIF
    if args.folder:
        folder_path = PathlibPath(args.folder)
        if folder_path.exists():
            print(f"为文件夹 {folder_path} 生成GIF")
            generate_gif_for_single_folder(folder_path)
        else:
            print(f"文件夹 {folder_path} 不存在!")
    else:
        # Constants
        flag = 'all_scenarios'  # default
        # 所有场景
        if flag == 'all_scenarios':
            dir_current_file = os.path.dirname(__file__)
            dir_current_file = PathlibPath(dir_current_file)
            for scenario_name in dir_current_file.iterdir():
                if scenario_name.suffix == '.py':
                    continue
                print(scenario_name.stem)
                generate_gif_for_single_folder(scenario_name)

        # 指定单个场景
        else:
            IMAGE_DIR_NAME = "image_B202_mission_0_0_1"
            # IMAGE_DIR_NAME = "image_jiangtong_intersection_2_8_4"
            # 自动计算路径
            BASE_DIR = PathlibPath(__file__).parent
            target_folder = BASE_DIR / IMAGE_DIR_NAME
            generate_gif_for_single_folder(target_folder)